using UnityEngine;
using System;
using System.Linq;
using System.Collections;
using System.Collections.Generic;
using DanielSig.UnityExtensions;
using DanielSig;

public class LearningAgent2 : MonoBehaviour, IAgent
{
	static float SigmoidDerivative(float value)
	{
		float temp = 1 + Mathf.Exp(value);
		return (2*Mathf.Exp(value))/(temp * temp);
	}
	static float Sigmoid01Derivative(float value)
	{
		float temp = 1 + Mathf.Exp(value);
		return (Mathf.Exp(value))/(temp * temp);
	}
	static float Sigmoid(float value)
	{
		return 2/(1+Mathf.Exp(-value))-1;
	}
	static float Sigmoid01(float value)
	{
		return 1/(1+Mathf.Exp(-value));
	}
	public enum LearningMethod
	{
		None,
		Observation,
		PunishmentAndReward
	}
	[System.Serializable]
	public class TeacherSettings
	{
		public LayerMask StaticObjectsMask;
		public string OpponentTag;
		public int PathFindingInterval = 20;
		public ProjectileController ProjectilePrefabController;
	}
	public class Neuron
	{
		public float Charge;
		public float ErrorDelta;
		public Neuron[] Parents;
		public float[] Weights;
		
		public Neuron()//input neuron
		{
			Parents = new Neuron[0];
			Weights = new float[0];
		}
		
		public Neuron(Neuron[] parents)//hidden and output neuron
		{
			Parents = parents;
			Weights = new float[parents.Length + 1];//+1 for bias
			for(int i = 0; i < Weights.Length; i++)
				Weights[i] = UnityEngine.Random.value * 2 - 1;
		}
		
		public Vector3 DrawPositionCache;
	}
	public static LinkedList<Neuron[]> Brain;//procedurally generated
	public int[] NeuronsPerLayer;//defined in unity
	[Range(0,1)]
	public float ExplorationAmount = 0.5F;
	[Range(0,1F)]
	public float LearningAmount = 0.1F;
	public LearningMethod Method = LearningMethod.Observation;
	public Vector3 DrawCenter;//defined in unity
	public Vector3 DrawNeuronStep = Vector3.right;//defined in unity
	public Vector3 DrawLayerStep = Vector3.back;//defined in unity
	public float DrawNeuronSize = 1;//defined in unity
	
	private ISensor[] _prevInput;//cache
	private Move _testMove;//cache
	private bool _waitingForFeedback = false;//should return Move.None?
	private float _curiosity = 0;
	private MasterAgent _teacher;
	public TeacherSettings Teacher;
	private static bool _restart = true;
	
	private static float _explorationAmount = -1;
	
	void Awake()
	{
		if(_explorationAmount >= 0)
			ExplorationAmount = _explorationAmount;
	}
	void OnDestroy()
	{
		_explorationAmount = ExplorationAmount;
	}
	
	public Move NextMove (ISensor[] input)
	{	
		if(_restart)
		{
			_restart = false;
			InitBrain(input);
		}
		if(Method == LearningMethod.PunishmentAndReward)
		{
			if(_waitingForFeedback)
				return _testMove;
			
			_curiosity += ExplorationAmount;
			if(_curiosity >= 1)
			{
				_curiosity--;
				_testMove = (Move)(int)(UnityEngine.Random.value * 15);
				_waitingForFeedback = true;
			}
		}
		
		Observe(input);
		Think();
		Move decision = Decide();
		if(Method == LearningMethod.PunishmentAndReward && _waitingForFeedback)
		{
			return _testMove;
		}
		else if(Method == LearningMethod.Observation)
		{
			Move teacherMove = GetTeacherDecision(input);
			ErrorCorrect(1, teacherMove);
			if(decision == Move.None || decision == (Move.TurnLeft | Move.TurnRight))
				decision = teacherMove;
		}
		
		return decision;
	}
	private Move GetTeacherDecision (ISensor[] input)
	{
		if(!_teacher)
		{
			_teacher = gameObject.RequestComponent<MasterAgent>();
			_teacher.StaticObjectsMask = Teacher.StaticObjectsMask;
			_teacher.OpponentTag = Teacher.OpponentTag;
			_teacher.PathFindingInterval = Teacher.PathFindingInterval;
			_teacher.ProjectilePrefabController = Teacher.ProjectilePrefabController;
		}
		
		return _teacher.NextMove(input);
	}
	public Move Decide ()
	{
		Move move = Move.None;
		Move current = (Move)1;
		Neuron[] output = Brain.Last.Value;
		
		foreach(Neuron n in output)
		{
			if(n.Charge > 0.9F)
			{
				move |= current;
			}
			current = (Move)((int)current << 1);
		}
		return move;
	}
	public void Think ()
	{
		LinkedListNode<Neuron[]> layer = Brain.First.Next;
		while(layer != null)//for each layer in brain besides input layer
		{
			foreach(Neuron n in layer.Value)
			{
				n.Charge = n.Weights[n.Parents.Length];//init to bias
				int i = 0;
				foreach(Neuron parent in n.Parents)
				{
					n.Charge += parent.Charge * n.Weights[i++];
				}
				n.Charge = Sigmoid01(n.Charge);//value will be between 0 and 1
			}
			layer = layer.Next;
		}
	}
	public void Observe (ISensor[] input)
	{
		_prevInput = input;//for rendering gizmos
		
		Neuron[] inputNeurons = Brain.First.Value;
		int neuronIndex = 0;
		foreach(ISensor sensor in input)
		{
			if(sensor.HasFixedDirection)//maps to 5 neurons
			{
				Color color = ColorizeSensation(sensor.Type);
				float dist = sensor.Values[0];
				inputNeurons[neuronIndex++].Charge = color.r;
				inputNeurons[neuronIndex++].Charge = color.g;
				inputNeurons[neuronIndex++].Charge = color.b;
				inputNeurons[neuronIndex++].Charge = color.a;
				inputNeurons[neuronIndex++].Charge = Sigmoid(dist);
			}
			else//maps to [sensor.Values.Length] neurons
			{
				foreach(float value in sensor.Values)
				{
					inputNeurons[neuronIndex++].Charge = Sigmoid(value);
				}
			}
		}
	}
	private Color ColorizeSensation(SensationType type)
	{
		Color c = Color.clear;
		switch (type)
		{
			case SensationType.Projectile: c = Color.white; break;
			case SensationType.Enemy: c = Color.blue; break;
			case SensationType.EnemyAmmo: c = Color.cyan.SetAlpha(0); break;
			case SensationType.EnemyHealth: c = Color.magenta.SetAlpha(0); break;
			case SensationType.None: c = Color.black; break;
			case SensationType.Obstacle: c = Color.clear; break;
			case SensationType.AmmoPack: c = Color.green; break;
			case SensationType.MyAmmo: c = Color.green.SetAlpha(0); break;
			case SensationType.HealthPack: c = Color.red; break;
			case SensationType.MyHealth: c = Color.red.SetAlpha(0); break;
			default: return c;
		}
		c.a = 1 - c.a;
		return c;
	}
	
	public void InitBrain (ISensor[] input)
	{
		//init
		Brain = new LinkedList<Neuron[]>();
		Neuron[] workingLayer;
		Neuron[] parentLayer;
		
		#region input layer
		int numInputs = input.Aggregate<ISensor,int>(0, (c,s) => c + (s.HasFixedDirection ? 5 : s.Values.Length));

		workingLayer = new Neuron[numInputs];
		for(int i = 0; i < numInputs; i++)
			workingLayer[i] = new Neuron();
		Brain.AddLast(workingLayer);
		#endregion
		
		#region hidden layer
		
		foreach(int size in NeuronsPerLayer)
		{
			parentLayer = workingLayer;
			workingLayer = new Neuron[size];
			for(int i = 0; i < size; i++)
				workingLayer[i] = new Neuron(parentLayer);
			Brain.AddLast(workingLayer);
		}
		#endregion
		
		#region output layer
		parentLayer = workingLayer;
		workingLayer = new Neuron[4];
		for(int i = 0; i < 4; i++)
			workingLayer[i] = new Neuron(parentLayer);
		Brain.AddLast(workingLayer);
		#endregion
	}

	private void ErrorCorrect(float amount, Move correctMove)
	{
		amount = Sigmoid(amount);
		amount *= LearningAmount;
		
		#region backpropagation
		
		int neuronIndex = 0;
		Neuron[] output = Brain.Last.Value;
		
		foreach(Neuron n in output)
		{
			int current = 1 << (3 - neuronIndex++);
			float correct = (((int)correctMove) & current) > 0 ? 1 : 0;
			float thought = n.Charge;
			n.ErrorDelta = thought - correct;
		}
		
		Action<Neuron,float> recursion;
		
		recursion = (n,correction) =>
		{
			int numberOfParents = n.Parents.Length;
			float correctValue = (n.Charge + correction) / n.Charge;
			bool notLastLayer = n.Parents[0].Parents.Length > 0;
			for(int i = 0; i < numberOfParents; i++)
			{
				n.Weights[i] = n.Weights[i] * (1-amount) + correctValue * amount;
				if(notLastLayer)
					recursion(n.Parents[i], correction * n.Weights[i].Abs());
			}
		};
		
		for(int i = 0; i < 4; i++)
			recursion(output[i], output[i].ErrorDelta);
		/*
		
		LinkedListNode<Neuron[]> layer = Brain.Last;
		while(layer.Previous != null)//for each layer in brain besides input layer, reverse order
		{
			int numParents = layer.Previous.Value.Length;
			int nIndex = 0;
			
			foreach(Neuron n in layer.Previous.Value)//for each neuron in layer
			{
				n.ErrorDelta = 0;
			}
			foreach(Neuron n in layer.Value)//for each neuron in layer
			{
				for(int i = 0; i < numParents; i++)//for each axiom connected to neuron
				{
					if(n.Weights[i] < 0.0001F)
						n.Weights[i] = 0.0001F;
					n.Parents[i].ErrorDelta += n.ErrorDelta * n.Weights[i];
				}
				nIndex++;
			}
			layer = layer.Previous;
		}
		layer = Brain.First.Next;
		while(layer != null)//for each layer in brain besides input layer
		{
			int numParents = layer.Previous.Value.Length;
			int nIndex = 0;
			foreach(Neuron n in layer.Value)//for each neuron in layer
			{
				//calculate new charge
				n.Charge = n.Weights[n.Parents.Length];//init to bias
				int i = 0;
				foreach(Neuron parent in n.Parents)
				{
					n.Charge += parent.Charge * n.Weights[i++];
				}
				n.Charge = Sigmoid01Derivative(n.Charge);//value will be between 0 and 1
				
				//adjust weights
				for(i = 0; i < numParents; i++)//for each axiom connected to neuron
				{
					n.Weights[i] = n.Weights[i] + amount * n.ErrorDelta * n.Charge * n.Parents[i].Charge;
				}
				n.Weights[numParents] = n.Weights[numParents] + amount * n.ErrorDelta * n.Charge;
				
				//recalculate charge
				n.Charge = n.Weights[n.Parents.Length];//init to bias
				i = 0;
				foreach(Neuron parent in n.Parents)
				{
					n.Charge += parent.Charge * n.Weights[i++];
				}
				n.Charge = Sigmoid01(n.Charge);
				
				nIndex++;
			}
			layer = layer.Next;
		}
		*/
		#endregion
	}
	void Restart()
	{
		_restart = true;
	}
	void Good(float amount)
	{
		if(!_waitingForFeedback)
			return;
		ErrorCorrect(amount, _testMove);
		_waitingForFeedback = false;
	}
	void Bad(float amount)
	{
		Good(-amount);
	}
	
	void OnDrawGizmos()
	{
		if(_prevInput == null || Brain == null)
			return;
		
		DrawInputs();
		DrawOtherLayers();
	}
	void DrawInputs()
	{
		Vector3 center = DrawCenter;
		Vector3 pos = center - DrawNeuronStep * ((_prevInput.Length / 2F) - 0.5F);
		
		Neuron[] inputNeurons = Brain.First.Value;
		int neuronIndex = 0;
		
		foreach(ISensor sensor in _prevInput)
		{
			if(sensor.HasFixedDirection)
			{
				Color color = ColorizeSensation(sensor.Type);
				Color c = color.SetAlpha(1);
				Gizmos.color = Color.Lerp(c, Color.gray, color.a * 0.5F).SetAlpha(1 / sensor.Values[0].Max(1));
				for(int i = 0; i < 5; i++)
				{
					inputNeurons[neuronIndex++].DrawPositionCache = pos;
				}
			}
			else
			{
				Gizmos.color = Color.Lerp(Color.white, Color.black, Sigmoid01(sensor.Values[0]));
				for(int i = 0; i < sensor.Values.Length; i++)
				{
					inputNeurons[neuronIndex++].DrawPositionCache = pos;
				}
			}
			Gizmos.DrawCube(pos, Vector3.one * DrawNeuronSize);
			pos += DrawNeuronStep;
		}
	}
	void DrawOtherLayers()
	{
		Vector3 center = DrawCenter;
		Vector3 pos;
		
		LinkedListNode<Neuron[]> layer = Brain.First.Next;
		while(layer != null)//for each layer in brain besides input layer
		{
			center += DrawLayerStep;
			pos = center - DrawNeuronStep * ((layer.Value.Length / 2F) - 0.5F);
			foreach(Neuron n in layer.Value)
			{
				//draw neuron
				Gizmos.color = Color.Lerp(Color.black, Color.white, Sigmoid01(n.Charge));
				Gizmos.DrawSphere(pos, DrawNeuronSize / 2);
				
				int i = 0;
				foreach(Neuron parent in n.Parents)
				{
					//draw axioms
					Gizmos.color = Color.Lerp(Color.black, Color.white, n.Parents[i].Charge * n.Weights[i]).SetAlpha(n.Weights[i++]);
					Gizmos.DrawLine(pos, parent.DrawPositionCache);
				}
				n.DrawPositionCache = pos;
				pos += DrawNeuronStep;
			}
			layer = layer.Next;
		}
	}
}

